#include "nm_string.h"

kstringw::kstringw()
{
	m_stringw = NULL;
	m_len = 0;
}

kstringw::kstringw( PWCHAR String )
{
	ULONG Size;

	Size = wcslen(String);
	m_stringw = Allocate(Size, FALSE);
	RtlStringCchCopyW(m_stringw, Size, String);

	m_len = Size;
}

kstringw::kstringw( PUNICODE_STRING UniString )
{
	ULONG Size;

	Size = (ULONG)UniString->MaximumLength;
	m_stringw = Allocate(Size, FALSE);
	RtlStringCchCopyNW(m_stringw, Size, UniString->Buffer, UniString->Length / sizeof(WCHAR));

	m_len = Size;
}

void kstringw::Release()
{
	if (m_stringw != NULL) {

		delete[] m_stringw;
	}

	m_len = 0;
}

PWCHAR kstringw::Allocate(ULONG Size, BOOLEAN SaveData)
{
	PWCHAR buf;

	if (0 == Size) {

		Size = 1;
	}

	if (m_len < Size) {
	
		m_len = Size;
		buf = new WCHAR[m_len + 1];
		RtlZeroMemory(m_stringw, (m_len + 1) * sizeof(WCHAR));
		if (SaveData) {

			if (NULL != m_stringw) {

				RtlStringCchCopyW(buf, Size, m_stringw);
			}
		}
		if (NULL != m_stringw) {
			
			delete[] m_stringw;
		}

		m_stringw = buf;
	}

	if (!SaveData) {
	
		RtlZeroMemory(m_stringw, (m_len + 1) * sizeof(WCHAR));
	}

	return m_stringw;
}

void kstringw::Append( WCHAR ch )
{
	DWORD len;
	PWCHAR buf;

	len = wcslen(m_stringw);

	buf = Allocate(len + 1, TRUE);

	buf[len] = ch;
	buf[len + 1] = L'\0';
}

void kstringw::Attach( PWCHAR String )
{
	if (NULL != m_stringw) {

		delete[] m_stringw;
	}

	m_stringw = String;
	m_len = wcslen(String);
}

int kstringw::Compare( PWCHAR String, BOOLEAN CaseInsensitive )
{
	if (CaseInsensitive) {
		
		return _wcsicmp(m_stringw, String);
	}
	else {

		return wcscmp(m_stringw, String);
	}
}

void kstringw::Copy( PWCHAR String, ULONG Cch )
{
	ULONG len;
	PWCHAR p;
	if (Cch == 0) {

		len = wcslen(String);
	}
	else {
		len = Cch;
	}
	
	p = Allocate(len, FALSE);

	if (0 == m_len) {
		Empty();
	}
	else {
		RtlStringCchCopyNW(p, len, String, len);
	}
}

PWCHAR kstringw::Detach()
{
	PWCHAR ret = m_stringw;
	m_stringw = NULL;
	m_len = 0;
	return ret;
}

void kstringw::Empty()
{
	if (NULL == m_stringw) {
		m_stringw = Allocate(1, FALSE);
	}
	m_stringw[0] = L'\0';
}

INT kstringw::Find( WCHAR ch, INT Start )
{
	PWCHAR p;

	if (NULL == m_stringw || (INT)wcslen(m_stringw) < Start) {
		return -1;
	}

	p = wcschr(m_stringw + Start, ch);
	if (NULL == p) {

		return -1;
	}

	return (INT)(p - m_stringw);
}

INT kstringw::Find( PWCHAR String, INT Start )
{
	PWCHAR p;

	if (NULL == m_stringw || (INT)wcslen(m_stringw) < Start) {
		return -1;
	}

	p = wcsstr(m_stringw + Start, String);
	if (NULL == p) {

		return -1;
	}

	return (INT)(p - m_stringw);
}

INT kstringw::Format( __in PWCHAR Formats, ... )
{
	WCHAR tmp[1024];
	va_list argList;
	INT ret;

	va_start(argList, Formats);
	ret = RtlStringCbPrintfW(tmp, 1024, Formats, argList);
	va_end(argList);

	Copy(tmp);
	return ret;
}

WCHAR kstringw::GetAt( __in ULONG idx )
{
	if (idx > (wcslen(m_stringw) - 1)) {

		idx = wcslen(m_stringw);
	}

	return m_stringw[idx];
}

INT kstringw::GetLength() const
{
	return wcslen(m_stringw);
}

PWCHAR kstringw::GetString() const
{
	return m_stringw;
}

PUNICODE_STRING kstringw::GetUnicodeString()
{
	RtlInitUnicodeString(&m_UniString, m_stringw);

	return &m_UniString;
}

BOOLEAN kstringw::IsEmpty() const
{
	return (NULL == m_stringw || L'\0' == *m_stringw);
}

void kstringw::Left( ULONG Cch )
{
	if (Cch < wcslen(m_stringw)) {
		
		m_stringw[Cch] = L'\0';
	}
}

void kstringw::Mid( ULONG Start, ULONG Cch )
{
	PWCHAR buf;

	buf = new WCHAR[Cch + 1];
	RtlZeroMemory(buf, (Cch + 1) * sizeof(WCHAR));

	RtlStringCchCopyNW(buf, Cch, m_stringw + Start, Cch);
	Copy(buf);

	delete[] buf;
}

void kstringw::MakeUpper()
{
	_wcsupr(m_stringw);
}

void kstringw::MakeLower()
{
	_wcslwr(m_stringw);
}

void kstringw::Right( ULONG Cch )
{
	PWCHAR buf;

	buf = new WCHAR[Cch + 1];
	RtlZeroMemory(buf, (Cch + 1) * sizeof(WCHAR));

	RtlStringCchCopyNW(buf, Cch, m_stringw + wcslen(m_stringw) - Cch, Cch);
	Copy(buf);

	delete[] buf;
}

INT kstringw::Replace( PWCHAR Old, PWCHAR New )
{
	ULONG StrLen = wcslen(m_stringw);
	ULONG OldLen = wcslen(Old);
	ULONG NewLen = wcslen(New);
	INT cnt = 0;
	PWCHAR Start = m_stringw;
	PWCHAR End = m_stringw + OldLen;
	PWCHAR Target;
	PWCHAR NewBuf;
	PWCHAR dst;
	PWCHAR src;
	PWCHAR p = NULL;
	INT CchCopy = 0;
	INT tmp;
	INT i;

	if (0 == StrLen) {
		return 0;
	}

	if (0 == OldLen) {
		return 0;
	}

	
	while (Start < End) {
		
		while ((Target = wcsstr(Start, Old)) != NULL) {
			cnt++;
			Start = Target + OldLen;
		}

		if (NULL == Target) {
			
			break;
		}
	}

	if (0 == cnt) {
		return 0;
	}

	StrLen += (NewLen - OldLen) * cnt;
	NewBuf = new WCHAR[StrLen];
	RtlZeroMemory(NewBuf, StrLen * sizeof(WCHAR));
	dst = NewBuf;
	src = m_stringw;
	tmp = StrLen;

	for (i = 0; i < cnt; i++) {

		p = wcsstr(src, Old);

		CchCopy = p - src;
		RtlStringCchCopyNW(dst, tmp, src, CchCopy);
		dst += CchCopy;
		tmp -= CchCopy;
		src = p;

		RtlStringCchCopyNW(dst, tmp, New, NewLen);
		dst += NewLen;
		tmp -= NewLen;
		src += OldLen;
	}

	RtlStringCchCopyW(dst, tmp, src);

	delete[] m_stringw;
	m_stringw = NewBuf;
	m_len = StrLen;
	return cnt;
}

INT kstringw::ReverseFind( WCHAR ch )
{
	PWCHAR p = wcsrchr(m_stringw, ch);
	if (NULL == p)
		return -1;
	return p - m_stringw;
}

WCHAR kstringw::SetAt( ULONG pos, WCHAR ch )
{
	ULONG len;
	WCHAR ret;
	if (NULL == m_stringw || L'\0' == *m_stringw)
		return -1;

	len = wcslen(m_stringw);
	if (pos >= len) {
		
		pos = len - 1;
	}
	ret = m_stringw[pos];
	m_stringw[pos] = ch;
	return ret;
}

kstringw::operator PWCHAR() const
{
	return m_stringw;
}

kstringw::operator PUNICODE_STRING()
{
	RtlInitUnicodeString(&m_UniString, m_stringw);

	return &m_UniString;
}

kstringw& kstringw::operator=( PWCHAR String )
{
	Copy(String);
	return *this;
}

kstringw& kstringw::operator=( PUNICODE_STRING UniString )
{
	Copy(UniString->Buffer, UniString->Length / sizeof(WCHAR));
	return *this;
}

kstringw& kstringw::operator=( const kstringw& str )
{
	if (this != &str)
		Copy((PWCHAR)str);
	return *this;
}

BOOLEAN kstringw::operator==( kstringw& str2 )
{
	if (this->Compare(str2.GetString()) == 0) {

		return TRUE;
	}

	return FALSE;
}

BOOLEAN kstringw::operator==( PWCHAR str2 )
{
	if (this->Compare(str2) == 0) {

		return TRUE;
	}

	return FALSE;
}


const kstringw& kstringw::operator+=( PUNICODE_STRING UniString )
{
	ULONG Len = GetLength();
	Allocate(Len + UniString->Length / sizeof(WCHAR), TRUE);
	RtlStringCchCopyNW(&m_stringw[Len], UniString->Length / sizeof(WCHAR), UniString->Buffer, UniString->Length / sizeof(WCHAR));
	return *this;
}

const kstringw& kstringw::operator+=( PWCHAR String )
{
	ULONG Len = GetLength();
	Allocate(Len + wcslen(String), TRUE);
	RtlStringCchCopyW(&m_stringw[Len], wcslen(String), String);
	return *this;
}

const kstringw& kstringw::operator+=( WCHAR ch )
{
	Append(ch);
	return *this;
}

const kstringw& kstringw::operator+=( const kstringw& str )
{
	ULONG Len = GetLength();
	Allocate(Len + str.GetLength(), TRUE);
	RtlStringCchCopyW(&m_stringw[Len], str.GetLength(), (PWCHAR)str);
	return *this;
}

WCHAR kstringw::operator[]( ULONG idx )
{
	return GetAt(idx);
}

